from .postprocess import *
import time

def mapillay_postprocess(result, conf, postprocess=False,
                     previous_segmentation_map=None):
    if conf is not None:
        seg_logits = result.seg_logits.data
        probs = torch.softmax(seg_logits, dim=0)
        max_probs, pred_labels = torch.max(probs, dim=0)
        uncertain_mask = max_probs < conf  # 低置信度区域掩码
        pred_labels[uncertain_mask] = -1  # 将低置信度像素标记为 -1（未知）
        result.pred_sem_seg.data = pred_labels

    if postprocess:
        t0 = time.time()
        ori_res = result.pred_sem_seg.data.to(torch.int8)
        # ori_res = result.pred_sem_seg.data.cpu()
        # print(f"\t转换cpu: {round((time.time() - t0) * 1000, 3)} ms")
        filter_t = time.time()
        post_res = filter_by_spatial_prior(
            ori_res,
            filter_rules={
                # 27: {"max_y_ratio":0.5},# 过滤位置偏下的“天空”类别
                # 17: {"max_y_ratio": 0.6},# 过滤位置偏下的“楼房”类别
                # 30: {"max_y_ratio": 0.8},# 过滤位置偏下的“植物”类别
                9: {"max_y_ratio": 0.5},  # 过滤位置偏下的“天空”类别
            },
            filled_label=-1
        )
        filter_end_t = time.time()
        print(f"\t过滤耗时: {round((filter_end_t-filter_t)*1000,3)} ms")
        # fill_result = fill_uncertain_pixels(
        #     morphology_result_cls, uncertain_label=-1)
        post_res = fill_uncertain_pixels_temporal(
            post_res,
            previous_segmentation_map=previous_segmentation_map.pred_sem_seg.data \
                if previous_segmentation_map is not None else None,
            uncertain_label=-1,
            default_fill_value=-1,
            allowed_previous_labels=[0, 1]  # 道路类
            # allowed_previous_labels=[7, 9, 10, 11, 13, 14, 15]  # 道路类
        )
        fill_end_t = time.time()
        print(f"\t填充耗时: {round((fill_end_t-filter_end_t) * 1000, 3)} ms")

        post_res = post_process_roads_morphology(
            post_res,
            road_labels=[0,1],
            # [7, 9, 10, 11, 13, 14, 15],
            dila_kernel_size=(9, 9),
            open_kernel_size=(9, 9),
            expect_labels_lst=[5, 6],    # 标线
            # expect_labels_lst=[8,23,24],
            min_component_size=1500,
            fill_unknown=True
        )
        morphology_end_t = time.time()
        print(f"\t形态学耗时: {round((morphology_end_t-fill_end_t) * 1000, 3)} ms")

        result.pred_sem_seg.data = post_res
    return result